#!/usr/bin/env python
# deviceio.relay.py
#
#  tcprelay wrapper for python
#
# pylint: disable=C0111, C0326, C0103, C0301

from __future__ import unicode_literals

import os
import re
import socket
import subprocess
import sys
import threading
import time
import fcntl
import traceback

import paramiko
from PyQt5.QtCore import pyqtSignal, QThread
import stat
from PyQt5.QtWidgets import QWidget

from conf.configuration import logger, replace_unconditionally, tool_path
from Widgets.NotificationBox import NoteBox


tcprelay_path = os.path.join(tool_path, 'tcprelay')


class TcprelayLock(object):

    def __init__(self, filename):
        self.filename = filename
        # This will create it if it does not exist already
        self.handle = open(filename, 'w')

    # Bitwise OR fcntl.LOCK_NB if you need a non-blocking lock
    def acquire(self):
        fcntl.flock(self.handle, fcntl.LOCK_EX)

    def release(self):
        fcntl.flock(self.handle, fcntl.LOCK_UN)

    def __del__(self):
        self.handle.close()


# logger = logging.getLogger(__name__)


class TcpRelayError(Exception):
    pass


class TcpRelay(object):
    def __init__(self, locationid):
        """
        Arguments:
            locationid: (integer,hex string)    The locationid of the device
            logger:     (logger logger)        Logger for class
        """
        assert (isinstance(locationid, int))

        if locationid == 0:
            logger.debug("Bad Location ID: {}".format(locationid))
            logger.debug("Attempting to find valid locationID")
            locationid = get_location_id_from_tcp_relay()

        self._relayproc = None
        self._relaypid = None
        self._locationid = locationid
        self._portoffset = None
        self._logfile = "/tmp/tcprelay.{0}.log"
        self._target = 'localhost'
        self._pwd = 'alpine'
        self._ssh_client = None
        self._lock = TcprelayLock("/tmp/py-tcprelay-open-port.lock")

    def __del__(self):
        self.close()

    @property
    def is_active(self):
        """
        Check if the tcprelay tunnel used by current object still valid
        """
        if not self._relaypid:
            return False

        self._lock.acquire()
        relaypid = None
        portoffset = None
        try:
            relaypid, portoffset = self._check_tcprelay()
        except AttributeError:
            logger.debug(
                "No active TCPRELAY tunnel on locationid - {0}"
                "".format(self.locationid_param))
        finally:
            self._lock.release()

        return (
                self._relaypid == relaypid and
                self._portoffset == portoffset
        )

    def open(self):
        """
        Will open a tcprelay connection or if one already exists it will
        piggyback on that
        """
        self._lock.acquire()
        try:
            self._relaypid, self._portoffset = self._check_tcprelay()
            logger.debug(
                "PIGGYBACK TCPRELAY"
                "PID: {0} PORT: {1}".format(self._relaypid,
                                            self._portoffset))
        except AttributeError:
            # TODO: tcprelays might want to close when test is over???
            self._portoffset = get_available_portoffset()
            command = "{2} --autoexit --portoffset {0} " \
                      "--locationid {1} rsync telnet " \
                      "ssh > /tmp/tcprelay.{1}.log 2>&1" \
                      " &".format(self._portoffset, self.locationid_param, repr(tcprelay_path))
            logger.debug("SPAWNING TCPRELAY - {0}".format(command))
            child = subprocess.Popen(["bash", "-c", command], close_fds=True)
            time.sleep(0.5)
            try:
                self._relaypid, self._portoffset = self._check_tcprelay()
            except AttributeError:
                logger.error(
                    "FAILED to SPAWN TCPRELAY - CMD {0} "
                    "OUTPUT: {1} ERROR: {2} RC: {3}".format(command,
                                                            child.stdout,
                                                            child.stderr,
                                                            child.returncode))
        finally:
            self._lock.release()

    def _check_tcprelay(self):
        """
        Check the existing tcp relay mapped to our locationid
        Returns:
            tcp_relaypid, portoffset
        Raises AttributeError if no match
        """
        check = 'ps -e -opid -ocommand | grep tcprelay | grep -v ' \
                'grep | grep {0}'.format(self.locationid_param)
        output, _ = subprocess.Popen(
            ["bash", "-c", check], stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            universal_newlines=True).communicate()
        regex = re.search(
            r"^\s*([0-9]+).* --portoffset ([0-9]+).*", output)
        match = regex.groups()
        relaypid = int(match[0])
        portoffset = int(match[1])
        logger.debug(
            "CHECK TCPRELAY - CMD: {0} "
            "OUTPUT: {1} PID: {2} PORT: {3}".format(check,
                                                    output,
                                                    relaypid,
                                                    portoffset))
        return relaypid, portoffset

    @property
    def locationid_param(self):
        return hex(self._locationid).rstrip('L')

    def portoffset(self):
        """
        get port offset
        Returns:
            (integer) the port offset
        """
        return self._portoffset

    def close(self):
        """
        abandons the current tcprelay
        """
        # logger.info(f"killing pid {self._relaypid}")
        # os.kill(self._relaypid, 9)
        self._relaypid = None
        self._portoffset = None

    def _rsync(self, rsync):
        """
        internal rsync via expect
        Arguments:
            rsync:  (string) the rysnc command
        """
        rsync_expect = "set timeout -1; spawn {rsync}; expect " \
                       "\"Password:\"; send \"{pwd}\n\"; expect eof"
        command = ["/usr/bin/expect", "-c",
                   rsync_expect.format(rsync=rsync, pwd=self._pwd)]
        logger.debug(command)
        proc = subprocess.Popen(command,
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        std, err = proc.communicate()
        if len(std):
            logger.info(std)
        if len(err):
            logger.error(err)
        #     return False
        return True

    def rsync_to(self, local, remote, flags='-av', user='root'):
        """
        rsync to device from host
        Arguments:
            local:  source path  - local host
            remote: destination path -remote device
            flags:  rsync flags
            user:   rsync username
        Returns:
            (Boolean) whether rsync was success
        """
        if self._portoffset:
            tail = "/" if local[-1] == "/" else ''
            local = os.path.abspath(os.path.expanduser(local)) + tail
            return self._rsync(
                'rsync {flags} {src} '
                '{usr}@{target}:{port}{dest}'.format(
                    flags=flags, src=local, usr=user, target=self._target,
                    port=(self._portoffset + 873), dest=remote))
        return False

    def rsync_from(self, remote, local, flags='-av', user='root'):
        """
        rsync to device from host
        Arguments:
            remote: source path -remote device
            local:  dest path  - local host
            flags:  rsync flags
            user:   rsync username
        Returns:
            (Boolean) whether rsync was success
        """
        if self._portoffset:
            tail = "/" if local[-1] == "/" else ''
            local = os.path.abspath(os.path.expanduser(local)) + tail
            return self._rsync(
                'rsync {flags} {usr}'
                '@{target}:{port}{src} {dest}'.format(
                    flags=flags, src=remote, usr=user, target=self._target,
                    port=(self._portoffset + 873), dest=local))
        return False

    def kill(self):
        """ Kill the tcprelay associated with this class. """
        subprocess.check_output(['kill', str(self._relaypid)])
        self.close()
        logger.info(f"tcprelay exited: {self.is_active}")

    @property
    def is_device_attached(self):
        return self._locationid in get_all_location_ids_from_tcp_relay()


def get_all_location_ids_from_tcp_relay():
    """
    use tcprelay to get location ID for all connected USB devices,
    Returns:
        a list of matched location id decimal integers
    """
    # tcprelay has a default 5s timeout, but enumerates quickly
    # Run in the background and kill after 0.1s to avoid this
    cmd = f"{repr(tcprelay_path)} --list & sleep 0.1; kill $!"
    output = subprocess.check_output(cmd, shell=True)
    location_ids = re.findall(r"Location:\s*([0-9A-Fa-f]+)", output.decode('utf8'))

    # convert hex string in location_ids to integer
    def hex_to_int(x): return int(x, 16)

    return list(map(hex_to_int, location_ids))


def get_location_id_from_tcp_relay():
    location_ids = get_all_location_ids_from_tcp_relay()
    if len(location_ids):
        return location_ids[0]
    else:
        logger.error(
            "Cannot find a valid locationID. "
            "Is device panicked or not connected?")
        return 0


def get_available_portoffset(target="localhost"):
    """
    scans host ports for an available portoffset
    Returns:
        integer or None if no port offset exists
    """
    target_ip = socket.gethostbyname(target)
    for portoffset in range(10000, 61000, 1000):
        i = portoffset + 873
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        result = sock.connect_ex((target_ip, i))
        sock.close()
        if result != 0:
            logger.debug("port open {0}".format(portoffset))
            return portoffset
    return None


class SSH(QThread):
    output_signal = pyqtSignal(str)
    notebox_signal = pyqtSignal(dict)

    def __init__(self, port=10000, *args, **kwargs):
        super(SSH, self).__init__(*args, **kwargs)
        if hasattr(self.parent(), 'note_box_'):
            self.notebox_signal.connect(self.parent().note_box_)
        self.host = '127.0.0.1'
        self.user = 'root'
        self.pwd = 'alpine'
        self.port = port + 22
        logger.info(f"Initializing SSH with port {self.port}")
        self.__transport = None
        self.connect()
        self.ssh = paramiko.SSHClient()
        self.ssh._transport = self.__transport
        self.shell = self.ssh.invoke_shell()
        self.task = None

    def run(self) -> None:
        if isinstance(self.task, dict):
            getattr(self, self.task.get('method'))(self.task.get('src'), self.task.get('dst'))
        else:
            while True:
                r = self.shell.recv(512)
                """有一部分字符需要无条件替换: \x07|\x08\x1b[K"""
                if r:
                    r = replace_unconditionally.sub(b'', r).decode()
                    self.output_signal.emit(r)
        logger.info(f"Thread run over.")

    def tcprelay(self):
        """open tcprelay"""
        logger.info(os.popen('pkill tcprelay').read())
        logger.info(os.popen('tcprelay --portoffset 10000 22').read())

    def thread_tcprelay(self):
        thr = threading.Thread(target=self.tcprelay)
        thr.setDaemon(True)
        thr.start()
        # thr.join()

    def connect(self, count_: int = 0):
        sock = (self.host, self.port)
        try:
            transport = paramiko.Transport(sock)
        except Exception as e:
            count_ += 1
            logger.error(e)
            logger.error('Please make sure your device is in os mode.')
            logger.error('Please make sure your device is in os mode.')
            if count_ >= 2:
                raise Exception(1, 'GB is not ready in OS mode')
            else:
                return self.connect(count_)

            # sys.exit(0)
        try:
            transport.connect(username=self.user, password=self.pwd)
        except Exception as e:
            logger.error(e)
            logger.error(
                'Connect failed, please make sure your device in os mode and try whether you can connect by nanokdp/nanocom')
            # sys.exit(0)
        self.__transport = transport

    def close(self):
        try:
            self.ssh.close()
            self.__transport.close()
        except AttributeError as e:
            logger.error(e, exc_info=1)

    def run_cmd(self, cmd):
        try:
            ssh = paramiko.SSHClient()
            ssh._transport = self.__transport
            # run cmd
            logger.info(f'Start run cmd {cmd}')
            stdin, stdout, stderr = ssh.exec_command(cmd)
            # get resoult
            r = stdout.read().decode('utf8').strip()
            # error msg
            err = stderr.read().decode('utf8').strip()
            logger.info(f'CMD - {repr(cmd)}\n'
                        f'Result - {repr(r)}\n'
                        f'Error - {repr(err)}')
            if r:
                return {'color': 'green', 'result': r}
            elif err:
                return {'color': 'red', 'result': err}
            else:
                return {'color': 'green', 'result': ''}
        except Exception as e:
            logger.error(e, exc_info=1)
            self.notebox_signal.emit({'m': 'warning', 'c': str(e)})

    def quit(self) -> None:
        logger.debug(f"Closing SSHClient.")
        self.ssh.close()
        return super(SSH, self).quit()

    # def execute(self, cmd):
    #     logger.info(f'Start run cmd {repr(cmd)}')
    #     stdin, stdout, stderr = self.ssh.exec_command(cmd)
    #     # get resoult
    #     r = stdout.read().decode('utf8').strip()
    #     # error msg
    #     err = stderr.read().decode('utf8').strip()
    #     logger.info(f'cmd {repr(cmd)} run over')
    #     if cmd.startswith('cd '):
    #         self.chg_cwd()
    #     if err:
    #         self.output_signal.emit(err)
    #         return err
    #     else:
    #         self.output_signal.emit(r)
    #         return r
    #
    # def execute_quiet(self, cmd):
    #     logger.info(f'Start run cmd {repr(cmd)}')
    #     stdin, stdout, stderr = self.ssh.exec_command(cmd)
    #     # get resoult
    #     r = stdout.read().decode('utf8').strip()
    #     # error msg
    #     err = stderr.read().decode('utf8').strip()
    #     logger.info(f'cmd {repr(cmd)} run over')
    #     if cmd.startswith('cd '):
    #         self.chg_cwd()
    #     if err:
    #         return err
    #     else:
    #         return r
    def execute(self, cmd):
        # cmd += '\n'
        logger.info(f'Start run cmd {repr(cmd)}')
        try:
            self.shell.send(cmd)
        except OSError:  # Socket is closed
            return
        if cmd == 'reboot':
            self.close()
        # time.sleep(1)
        # r = self.shell.recv(1024).decode()
        # # get resoult
        # logger.info(f'cmd {repr(cmd)} run over')
        # self.output_signal.emit(r)
        # return r

    def execute_quiet(self, cmd):
        # cmd += '\n'
        logger.info(f'Start run cmd {repr(cmd)}')
        self.shell.send(cmd)
        # time.sleep(1)
        # r = self.shell.recv(1024).decode()
        # # get resoult
        # logger.info(f'cmd {repr(cmd)} run over')
        # return r

    def getf(self, src, dst, sftp=None):
        try:
            if sftp:
                sftp = sftp  # 这是getd 调用getf的时候，会传递sftp对象过来
                logger.debug(f"Getting file from {src} to {dst}")
                sftp.get(src, dst)
            else:
                sftp = paramiko.SFTPClient.from_transport(self.__transport)
                if os.path.basename(src) != os.path.basename(dst):
                    dst = os.path.join(dst, os.path.basename(src))
                logger.debug(f"Getting file from {src} to {dst}")
                sftp.get(src, dst)
                self.notebox_signal.emit({'c': f"Already copy {src} to {dst}", 'm': 'info'})
        except Exception as e:
            logger.error(e, exc_info=1)
            self.notebox_signal.emit({'m': 'warning', 'c': f'Failed to get file from {src} to {dst}'})

    def getd(self, src, dst):
        if os.path.basename(src) != os.path.basename(dst):
            dst = os.path.join(dst, os.path.basename(src))
            os.mkdir(dst) if not os.path.exists(dst) else ''  # 如果basename不一样，则创建文件夹；否则只copy子文件进来

        try:
            sftp = paramiko.SFTPClient.from_transport(self.__transport)
            logger.debug(f"Getting files from {src} to {dst}")
            all_files = self._get_all_files_in_remote_dir(sftp, src)
            assert all_files, f'Cannot find path {src}'
            for file in all_files:

                local_filename = file.replace(src, dst)
                local_filepath = os.path.dirname(local_filename)

                if not os.path.exists(local_filepath):
                    os.makedirs(local_filepath)

                self.getf(file, local_filename, sftp)
            else:
                # self.note_box.info(f"Already copy {os.path.basename(src)} to {dst}")
                self.notebox_signal.emit({'c': f"Already copy {os.path.basename(src)} to {dst}", 'm': 'info'})
        except Exception as e:
            logger.error(e, exc_info=1)
            # self.note_box.warning('ssh get dir from master failed.')
            self.notebox_signal.emit({'m': 'warning', 'c': f'{e}'})

    def putf(self, src, dst, sftp=None):
        try:
            if sftp:
                sftp = sftp
                logger.debug(f"Putting file from {src} to {dst}")
                sftp.put(src, dst)
            else:
                sftp = paramiko.SFTPClient.from_transport(self.__transport)
                if os.path.basename(src) != os.path.basename(dst):
                    dst = os.path.join(dst, os.path.basename(src))
                logger.debug(f"Putting file from {src} to {dst}")
                sftp.put(src, dst)
                self.notebox_signal.emit({'c': f"Already copy {src} to {dst}", 'm': 'info'})
        except Exception as e:
            logger.error(e, exc_info=1)
            self.notebox_signal.emit({'m': 'warning', 'c': f'Failed to get file from {src} to {dst}'})

    def putd(self, src, dst):
        if os.path.basename(src) != os.path.basename(dst):
            dst = os.path.join(dst, os.path.basename(src))

        logger.debug(f"Putting files from {src} to {dst}")
        try:
            sftp = paramiko.SFTPClient.from_transport(self.__transport)

            if dst[-1] == "/":
                dst = dst[0:-1]

            all_files = self._get_all_files_in_local_dir(src)
            for file in all_files:

                remote_filename = file.replace(src, dst)
                remote_path = os.path.dirname(remote_filename)

                try:
                    sftp.stat(remote_path)
                except:
                    # os.popen('mkdir -p %s' % remote_path)
                    self.run_cmd('mkdir -p %s' % remote_path)  # 使用这个远程执行命令

                self.putf(file, remote_filename, sftp)
            else:
                # self.note_box.info(f"Already copy {os.path.basename(src)} to {dst}")
                self.notebox_signal.emit({'m': 'info', 'c': f"Already copy {os.path.basename(src)} to {dst}"})

        except:
            self.notebox_signal.emit({'c':'ssh get dir from master failed.', 'm':'warning'})
            logger.error(traceback.format_exc())

    def lstate(self, path):
        sftp = paramiko.SFTPClient.from_transport(self.__transport)
        return sftp.stat(path)

    def _get_all_files_in_remote_dir(self, sftp=None, remote_dir: str = ''):
        sftp = sftp or paramiko.SFTPClient.from_transport(self.__transport)

        all_files = list()
        if remote_dir[-1] == '/':
            remote_dir = remote_dir[0:-1]
        try:
            files = sftp.listdir_attr(remote_dir)
            for file in files:
                filename = remote_dir + '/' + file.filename

                if stat.S_ISDIR(file.st_mode):  # 如果是文件夹的话递归处理
                    all_files.extend(self._get_all_files_in_remote_dir(sftp, filename))
                else:
                    all_files.append(filename)
        except FileNotFoundError as e:
            logger.error(e)
            # self.note_box.warning(f"No such file: {remote_dir}")
            self.notebox_signal.emit({'m': 'warning', 'c': f"No such file: {remote_dir}"})

        return all_files

    # 递归遍历本地服务器指定目录下的所有文件
    def _get_all_files_in_local_dir(self, local_dir: str = ''):
        all_files = list()

        for root, dirs, files in os.walk(local_dir, topdown=True):
            for file in files:
                filename = os.path.join(root, file)
                all_files.append(filename)

        return all_files

    def __del__(self):
        self.close()


if __name__ == '__main__':
    ssh = SSH()
    src = ''
    dst = ''
    ssh.getf(src, dst)
